#!/usr/bin/env python3
###
# @Author: Logan.Li
# @Gitee: https://gitee.com/attacker
# @email: admin@attacker.club
# @Date: 2025-03-14 09:00:30
# @LastEditTime: 2025-03-14 09:00:35
# @Description: 阿里云负载均衡SLB监控模块
###

import os
import json
import logging
from typing import Dict, List
from aliyunsdkslb.request.v20140515 import DescribeLoadBalancersRequest
from ..core import AliyunMonitorBase

class SLBMonitor(AliyunMonitorBase):
    """负载均衡监控"""
    
    def __init__(self, config: Dict):
        super().__init__(config)
        self.namespace = "acs_slb_dashboard"
        
    def list_load_balancers(self) -> List[Dict]:
        """获取负载均衡实例列表"""
        try:
            request = DescribeLoadBalancersRequest.DescribeLoadBalancersRequest()
            response = self.client.do_action_with_exception(request)
            return json.loads(response).get('LoadBalancers', {}).get('LoadBalancer', [])
        except Exception as e:
            logging.error(f"获取负载均衡实例列表失败: {e}")
            return []
            
    def setup_instance_monitoring(self, instance_id: str, instance_name: str):
        """设置实例监控
        
        Args:
            instance_id: 实例ID
            instance_name: 实例名称
        """
        dimensions = [{"instanceId": instance_id}]
        
        # QPS监控
        self.create_alarm_rule(
            rule_name=f"{instance_name}-QPS告警",
            namespace=self.namespace,
            metric_name="Qps",
            dimensions=dimensions,
            threshold=float(self.config['thresholds']['qps']),
            contact_groups=self.config['contact_groups']
        )
        
        # 活跃连接数监控
        self.create_alarm_rule(
            rule_name=f"{instance_name}-活跃连接数告警",
            namespace=self.namespace,
            metric_name="ActiveConnection",
            dimensions=dimensions,
            threshold=float(self.config['thresholds']['active_connection']),
            contact_groups=self.config['contact_groups']
        )
        
        # 每秒新建连接数监控
        self.create_alarm_rule(
            rule_name=f"{instance_name}-新建连接数告警",
            namespace=self.namespace,
            metric_name="NewConnection",
            dimensions=dimensions,
            threshold=float(self.config['thresholds']['new_connection']),
            contact_groups=self.config['contact_groups']
        )
        
        # 流入带宽监控
        self.create_alarm_rule(
            rule_name=f"{instance_name}-流入带宽告警",
            namespace=self.namespace,
            metric_name="TrafficRXNew",
            dimensions=dimensions,
            threshold=float(self.config['thresholds']['traffic_in']),
            contact_groups=self.config['contact_groups']
        )
        
        # 流出带宽监控
        self.create_alarm_rule(
            rule_name=f"{instance_name}-流出带宽告警",
            namespace=self.namespace,
            metric_name="TrafficTXNew",
            dimensions=dimensions,
            threshold=float(self.config['thresholds']['traffic_out']),
            contact_groups=self.config['contact_groups']
        )
        
        # 非健康后端服务器数监控
        self.create_alarm_rule(
            rule_name=f"{instance_name}-非健康后端服务器告警",
            namespace=self.namespace,
            metric_name="UnhealthyServerCount",
            dimensions=dimensions,
            threshold=float(self.config['thresholds']['unhealthy_server']),
            contact_groups=self.config['contact_groups']
        )
        
    def get_instance_metrics(self, instance_id: str) -> Dict:
        """获取实例的监控指标
        
        Args:
            instance_id: 实例ID
        """
        dimensions = [{"instanceId": instance_id}]
        metrics = {}
        
        metric_list = [
            ("Qps", "qps"),
            ("ActiveConnection", "active_connection"),
            ("NewConnection", "new_connection"),
            ("TrafficRXNew", "traffic_in"),
            ("TrafficTXNew", "traffic_out"),
            ("UnhealthyServerCount", "unhealthy_server")
        ]
        
        for metric_name, key in metric_list:
            data = self.get_metric_data(
                namespace=self.namespace,
                metric_name=metric_name,
                dimensions=dimensions
            )
            metrics[key] = data
            
        return metrics
        
    def monitor_all_instances(self):
        """监控所有负载均衡实例"""
        instances = self.list_load_balancers()
        for instance in instances:
            instance_id = instance['LoadBalancerId']
            instance_name = instance.get('LoadBalancerName', instance_id)
            self.setup_instance_monitoring(instance_id, instance_name)
            logging.info(f"已设置实例监控: {instance_name}")
            
    def get_listener_metrics(self, instance_id: str, port: int, protocol: str) -> Dict:
        """获取监听端口的指标
        
        Args:
            instance_id: 实例ID
            port: 端口号
            protocol: 协议（http/https/tcp/udp）
        """
        dimensions = [
            {"instanceId": instance_id},
            {"port": str(port)},
            {"protocol": protocol.lower()}
        ]
        
        metrics = {}
        
        # 状态码统计
        if protocol.lower() in ['http', 'https']:
            for status_code in ['2xx', '3xx', '4xx', '5xx']:
                data = self.get_metric_data(
                    namespace=self.namespace,
                    metric_name=f"Status{status_code}Count",
                    dimensions=dimensions
                )
                metrics[f'status_{status_code}'] = data
                
        # 响应时间
        metrics['response_time'] = self.get_metric_data(
            namespace=self.namespace,
            metric_name="UpstreamResponseTime",
            dimensions=dimensions
        )
        
        return metrics

# 使用示例
if __name__ == "__main__":
    # 配置示例
    config = {
        "access_key": "your_access_key",
        "secret_key": "your_secret_key",
        "region": "cn-hangzhou",
        "thresholds": {
            "qps": 1000,
            "active_connection": 5000,
            "new_connection": 500,
            "traffic_in": 100,  # Mbps
            "traffic_out": 100,  # Mbps
            "unhealthy_server": 1
        },
        "contact_groups": ["default"],
        "notification": {
            "dingtalk": {
                "enabled": True,
                "webhook": "your_webhook_url",
                "secret": "your_secret"
            },
            "lark": {
                "enabled": True,
                "webhook": "your_webhook_url"
            },
            "telegram": {
                "enabled": True,
                "bot_token": "your_bot_token",
                "chat_id": "your_chat_id"
            }
        }
    }
    
    monitor = SLBMonitor(config)
    
    # 监控所有实例
    monitor.monitor_all_instances()
    
    # 获取特定实例的指标
    metrics = monitor.get_instance_metrics("your_lb_id")
    print(json.dumps(metrics, indent=2))
    
    # 获取特定监听端口的指标
    listener_metrics = monitor.get_listener_metrics(
        instance_id="your_lb_id",
        port=80,
        protocol="http"
    )
    print("监听端口指标:", json.dumps(listener_metrics, indent=2))
